package net.sudot.commons.security;


import net.sudot.commons.json.JsonUtils;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.time.Duration;
import java.util.Base64;

/**
 * Json Web Token
 *
 * @author tangjialin on 2018-04-07.
 */
public class Jwt {
    private static final String ALGORITHM = "HmacSHA256";
    /**
     * 初始化head部分的数据为
     * {
     * "alg":"HS256",
     * "type":"JWT"
     * }
     */
    private static final String HEADER = Base64.getUrlEncoder().encodeToString("{\"alg\":\"HS256\",\"typ\":\"JWT\"}".getBytes(StandardCharsets.UTF_8));

    /**
     * 创建Token
     *
     * @param payload Token载体
     * @param secret  密钥
     * @return 返回Token
     */
    public static String create(Payload payload, String secret) {
        String content = Base64.getUrlEncoder().encodeToString(JsonUtils.toJsonBytes(payload));
        String signature = sign(HEADER + content, secret);
        StringBuilder builder = new StringBuilder(3 + HEADER.length() + content.length() + signature.length())
                .append(HEADER).append(".").append(content).append(".").append(signature);
        return builder.toString();
    }

    /**
     * 验证Token并获得载体信息
     *
     * @param token  Token
     * @param secret 密钥
     * @return 返回载体信息
     */
    public static Payload verify(String token, String secret) {
        if (token == null || token.isEmpty()) { return null; }
        int headerLastIndex = token.indexOf(".");
        int payloadLastIndex = token.lastIndexOf(".");
        if (headerLastIndex <= 0 || headerLastIndex >= payloadLastIndex) { return null; }
        String encodeHeader = token.substring(0, headerLastIndex);
        if (!HEADER.equals(encodeHeader)) { return null; }
        String encodePayload = token.substring(headerLastIndex + 1, payloadLastIndex);
        String encodeSignature = token.substring(payloadLastIndex + 1);
        String signature = sign(encodeHeader + encodePayload, secret);
        if (!encodeSignature.equals(signature)) { return null; }
        Payload payload = JsonUtils.toObject(java.util.Base64.getUrlDecoder().decode(encodePayload), Payload.class);
        if (payload.getExp() < System.currentTimeMillis()) {
            return null;
        }
        return payload;
    }

    /**
     * 签名
     *
     * @param string 需要签名的字符串
     * @param secret 前面密钥
     * @return 返回签名后的字符串
     */
    private static String sign(String string, String secret) {
        try {
            Mac mac = Mac.getInstance(ALGORITHM);
            mac.init(new SecretKeySpec(secret.getBytes(StandardCharsets.UTF_8), ALGORITHM));
            byte[] signatureBytes = mac.doFinal(string.getBytes(StandardCharsets.UTF_8));
            return Base64.getUrlEncoder().encodeToString(signatureBytes);
        } catch (NoSuchAlgorithmException | InvalidKeyException e) {
        }
        return null;
    }

    /**
     * 信息载体
     *
     * @author tangjialin on 2018-04-07.
     */
    public static class Payload implements Serializable {
        public static final Duration DEFAULT_EXP = Duration.ofHours(2);
        private String sub;
        private Long iat;
        private Long exp;

        public Payload() {
        }

        public Payload(String sub) {
            this(sub, DEFAULT_EXP);
        }

        public Payload(String sub, Duration expirationTime) {
            this.sub = sub;
            iat = System.currentTimeMillis();
            exp = iat + expirationTime.toMillis();
        }

        public String getSub() {
            return sub;
        }

        public Payload setSub(String sub) {
            this.sub = sub;
            return this;
        }

        public Long getIat() {
            return iat;
        }

        public Payload setIat(Long iat) {
            this.iat = iat;
            return this;
        }

        public Long getExp() {
            return exp;
        }

        public Payload setExp(Long exp) {
            this.exp = exp;
            return this;
        }
    }
}
